{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-07-03T06:38:41.548501Z",
     "start_time": "2024-07-03T06:38:40.364819Z"
    }
   },
   "source": [
    "import pyiqa\n",
    "from PIL import Image\n",
    "from torchvision.transforms import ToTensor\n",
    "import glob\n",
    "import os\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "niqe = pyiqa.create_metric(\"niqe\")\n",
    "nima = pyiqa.create_metric(\"nima\")\n",
    "brisque = pyiqa.create_metric(\"brisque\")\n",
    "psnr = pyiqa.create_metric(\"psnr\")\n",
    "ssim = pyiqa.create_metric(\"ssim\")\n",
    "lpips = pyiqa.create_metric(\"lpips\")\n",
    "\n",
    "IMG_EXTENSIONS = [\n",
    "    '.jpg', '.JPG', '.jpeg', '.JPEG',\n",
    "    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',\n",
    "    '.tif', '.TIF', '.tiff', '.TIFF', '.npy'\n",
    "]\n",
    "order_list = [\n",
    "    \"hazy\",\n",
    "    \"DCP\",\n",
    "    \"AOD\",\n",
    "    \"GCA\",\n",
    "    \"FFA\",\n",
    "    \"TNN\",\n",
    "    \"Dehamer\",\n",
    "    \"FogRemoval\",\n",
    "    \"ITB\",\n",
    "    \"SCA\",\n",
    "    \"ours\",\n",
    "    \"clean\"\n",
    "]\n",
    "\n",
    "def custom_order(x):\n",
    "    # 定义自定义排序规则\n",
    "    if \"hazy\" in x:\n",
    "        return 1\n",
    "    if \"DCP-DENSE-HAZE\" in x:\n",
    "        return 2\n",
    "    if \"AOD\" in x:\n",
    "        return 3\n",
    "    if \"GCA\" in x:\n",
    "        return 4\n",
    "    if \"FFA\" in x:\n",
    "        return 5\n",
    "    if \"TNN\" in x:\n",
    "        return 6\n",
    "    if \"Dehamer\" in x:\n",
    "        return 7\n",
    "    if \"FogRemoval\" in x:\n",
    "        return 8\n",
    "    if \"ITB\" in x:\n",
    "        return 9\n",
    "    if \"SCA\" in x:\n",
    "        return 10\n",
    "    if \"clean\" in x:\n",
    "        return 12\n",
    "    return 11\n",
    "\n",
    "def custom_order_for_img(x):\n",
    "    _, filename = os.path.split(x)\n",
    "    if \"_\" in filename:\n",
    "        return int(filename.split(\"_\")[0])\n",
    "    else:\n",
    "        return int(filename.split(\".\")[0])\n",
    "\n",
    "def get_image_list(root_path, dataset_name, clean_path, haze_path):\n",
    "    exp_dir = []\n",
    "    for root, dirs, files in os.walk(root_path):\n",
    "        exp_dir = dirs\n",
    "        break\n",
    "\n",
    "    exp_dir = [os.path.join(root_path, x) for x in exp_dir if dataset_name in x]\n",
    "    exp_dir.append(os.path.abspath(haze_path))\n",
    "    exp_dir.append(os.path.abspath(clean_path))\n",
    "\n",
    "    exp_dir = sorted(exp_dir, key=custom_order)\n",
    "    exp_imgs = []\n",
    "    for _, x in enumerate(exp_dir):\n",
    "        temp_list = glob.glob(x + \"/**/*.*\", recursive=True)\n",
    "        temp_list = [x for x in temp_list if os.path.splitext(x)[-1] in IMG_EXTENSIONS]\n",
    "        exp_imgs.append(sorted(temp_list, key=custom_order_for_img))\n",
    "    assert len(exp_imgs) == 12, \"Number of experiments should be 12: \" + str(len(exp_dir))\n",
    "    return exp_imgs\n",
    "\n",
    "def generate_comparison_images(root_path, dataset_name, clean_path, haze_path):\n",
    "    result_path = os.path.abspath(\"../compare/\" + dataset_name)\n",
    "    if not os.path.exists(result_path):\n",
    "        os.makedirs(result_path)\n",
    "    exp_imgs = get_image_list(root_path, dataset_name, clean_path, haze_path)\n",
    "\n",
    "    for count in range(len(exp_imgs[0])):\n",
    "        plt.rcParams['figure.dpi'] = 600\n",
    "        fig, axs = plt.subplots(1, 12)\n",
    "        # clean_img = ToTensor()(Image.open(exp_imgs[0][count]).convert('RGB'))[None, ::]\n",
    "        for i in range(len(exp_imgs)):\n",
    "            # current_img = ToTensor()(Image.open(exp_imgs[i][count]).convert('RGB'))[None, ::]\n",
    "            # title = \"\"\"\n",
    "            # {:s}\\n\n",
    "            # NIQE: {:.2f}\\n\n",
    "            # NIMA: {:.2f}\\n\n",
    "            # BRISQUE: {:.2f}\\n\n",
    "            # PSNR: {:.2f}\\n\n",
    "            # SSIM: {:.2f}\\n\n",
    "            # LPIPS: {:.2f}\n",
    "            # \"\"\".format(\n",
    "            #     order_list[i],\n",
    "            #     niqe(current_img).item(),\n",
    "            #     nima(current_img).item(),\n",
    "            #     brisque(current_img).item(),\n",
    "            #     psnr(clean_img, current_img).item(),\n",
    "            #     ssim(clean_img, current_img).item(),\n",
    "            #     lpips(clean_img, current_img).item()\n",
    "            # )\n",
    "            # \n",
    "            axs[i].imshow(plt.imread(exp_imgs[i][count]))\n",
    "            axs[i].set_title(order_list[i], fontproperties={'size': 4})\n",
    "            axs[i].axis('off')\n",
    "        plt.tight_layout()\n",
    "        plt.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0.01, hspace=0)\n",
    "        plt.savefig(os.path.join(result_path, 'COMPARE-' + str(count) +'.png'), bbox_inches='tight', pad_inches=0)\n",
    "        # plt.show()\n",
    "        plt.close()"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading pretrained model NIMA from /home/earthyzinc/.cache/torch/hub/checkpoints/NIMA_VGG16_ava-dc4e8265.pth\n",
      "Loading pretrained model LPIPS from /home/earthyzinc/.cache/torch/hub/checkpoints/LPIPS_v0.1_alex-df73285e.pth\n"
     ]
    }
   ],
   "execution_count": 14
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-03T05:31:52.769045Z",
     "start_time": "2024-07-03T05:29:56.244959Z"
    }
   },
   "cell_type": "code",
   "source": [
    "generate_comparison_images(\"/mnt/e/DeepLearningCopies/2023/RIDCP/results\",\n",
    "                           \"DENSE-HAZE\",\n",
    "                           \"/mnt/d/DeepLearning/dataset/Dense-Haze/clean\",\n",
    "                           \"/mnt/d/DeepLearning/dataset/Dense-Haze/hazy\")"
   ],
   "id": "b8ed42c723484dbe",
   "outputs": [],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-03T06:41:21.111543Z",
     "start_time": "2024-07-03T06:38:44.019435Z"
    }
   },
   "cell_type": "code",
   "source": [
    "generate_comparison_images(\"/mnt/e/DeepLearningCopies/2023/RIDCP/results\",\n",
    "                           \"NH-HAZE-20\",\n",
    "                           \"/mnt/d/DeepLearning/dataset/NH-HAZE-2020/clean\",\n",
    "                           \"/mnt/d/DeepLearning/dataset/NH-HAZE-2020/hazy\")"
   ],
   "id": "f36d92b90c96ecbc",
   "outputs": [],
   "execution_count": 15
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "generate_comparison_images(\"/mnt/e/DeepLearningCopies/2023/RIDCP/results\",\n",
    "                           \"NH-HAZE-21\",\n",
    "                           \"/mnt/d/DeepLearning/dataset/NH-HAZE-2021/clean\",\n",
    "                           \"/mnt/d/DeepLearning/dataset/NH-HAZE-2021/hazy\")"
   ],
   "id": "68208b0e92abab7f"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "generate_comparison_images(\"/mnt/e/DeepLearningCopies/2023/RIDCP/results\",\n",
    "                           \"NH-HAZE-23\",\n",
    "                           \"/mnt/d/DeepLearning/dataset/NH-HAZE-2023/clean\",\n",
    "                           \"/mnt/d/DeepLearning/dataset/NH-HAZE-2023/hazy\")"
   ],
   "id": "119c32cbdac794e6"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "generate_comparison_images(\"/mnt/e/DeepLearningCopies/2023/RIDCP/results\",\n",
    "                           \"O-HAZE\",\n",
    "                           \"/mnt/d/DeepLearning/dataset/O-HAZE/clean\",\n",
    "                           \"/mnt/d/DeepLearning/dataset/O-HAZE/hazy\")"
   ],
   "id": "8cf9a2495b5994ef"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "generate_comparison_images(\"/mnt/e/DeepLearningCopies/2023/RIDCP/results\",\n",
    "                           \"I-HAZE\",\n",
    "                           \"/mnt/d/DeepLearning/dataset/I-HAZE/clean\",\n",
    "                           \"/mnt/d/DeepLearning/dataset/I-HAZE/hazy\")"
   ],
   "id": "b7d6d0b4cae669a0"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
